[FlyDSL] add gather-gemm for a8w8 - #5207

Adds a FlyDSL implementation of DeepSeek MLA’s fused “gather KV + B-proj GEMM + ROPE copy” path for gfx950 (CDNA4) and a targeted test suite to validate correctness against Torch and Triton references.

Changes:

• Introduces an 8-wave FlyDSL kernel for gathered A (KV cache) × fp8 weight GEMM with optional per-row vs 128×128 block scaling, plus a ROPE copy epilogue.
• Adds a Python FlyDSL backend wrapper that validates inputs, compiles/caches the kernel, and launches it via the FlyDSL tensor shim.
• Adds extensive pytest coverage and exports the new API from aiter.ops.flydsl.

• Files reviewed: 4/4 changed files
• Comments generated: 2
• Review effort level: Lite

Suppressed comments (2)

aiter/ops/flydsl/kernels/gather_kv_b_proj_8wave.py:295

• _rope_copy iterates ROWS_PER_PASS (256) rows per pass, but for BLOCK_M values that are not a multiple of 256 (e.g. BLOCK_M=384) the second pass will process rows 384..511 and write rope for rows that belong to the next workgroup tile. This introduces cross-workgroup write races and can cause nondeterminism even if the computed values happen to match.

    for _p in range_constexpr(n_passes):
        row_local = tid // 2 + _p * ROWS_PER_PASS
        fx.copy(atom32, fx.slice(kvi_div, (None, m_base + row_local)), ireg)
        idx = Vec(fx.memref_load_vec(ireg))[0]
        src_base = idx * KV_ROW_ELEMS + KV_C_DIM + half * 32

aiter/ops/flydsl/gather_kv_b_proj.py:205

• _arch_of(k_buffer.device.index) will raise a non-ValueError exception if k_buffer is accidentally passed as a CPU tensor (device.index is None), which makes the failure mode inconsistent with the rest of the input validation.

    arch = _arch_of(k_buffer.device.index)

• Files reviewed: 4/4 changed files
• Comments generated: 0 new
• Review effort level: Lite

Suppressed comments (6)

Previously missed (2) — in code that hasn't changed since the last review.

aiter/ops/flydsl/gather_kv_b_proj.py:162

• The kv_proj_scale parameter comment only mentions per-row scales, but the function also accepts 128x128 block scales (2-D). Updating the signature comment avoids confusion for API users.

This issue also appears on line 191 of the same file.
op_tests/flydsl_tests/test_flydsl_gather_kv_b_proj.py:8

• The module docstring says this backend covers only per-output-row weight scales, but the tests in this file also cover 128x128 block scales. Updating the docstring keeps the stated supported configuration accurate.

aiter/ops/flydsl/gather_kv_b_proj.py:195

• The error message for kv_proj_scale=None says the backend is "per-output-row scale only", but this backend also supports 128x128 block-scale weights (see the kv_proj_scale shape validation later in the function). This message is misleading for callers/debugging.

    if kv_proj_scale is None:
        raise ValueError(
            "[FlyDSL gather_kv_b_proj] an unquantized weight (kv_proj_scale=None) "
            "is not supported; this backend is fp8 x per-output-row scale only. "
        )

aiter/ops/flydsl/gather_kv_b_proj.py:255

• k_scale is documented as shape [1], but the implementation will silently ignore extra elements (the kernel loads only element 0). Add a validation so incorrect shapes fail loudly instead of mis-scaling outputs.

    if m_rows == 0:
        return

aiter/ops/flydsl/gather_kv_b_proj.py:260

• kv_indices is assumed to be int32 by the kernel (it uses 32-bit buffer loads and Int32 arithmetic). If a caller passes int64, the kernel will reinterpret the bytes and produce incorrect gathers; validate the dtype explicitly.

    if kv_indices.numel() < m_rows:
        raise ValueError(
            f"[FlyDSL gather_kv_b_proj] kv_indices has {kv_indices.numel()} entries, "
            f"need at least num_tokens={m_rows}"
        )

aiter/ops/flydsl/gather_kv_b_proj.py:227

• k_prefix/v_prefix are flattened with .view(-1) for the kernel call; if either output is non-contiguous this raises a low-level RuntimeError. It’s better to validate contiguity and raise a clear ValueError before launching.

    if k_prefix.dtype != torch.bfloat16 or v_prefix.dtype != torch.bfloat16:
        raise ValueError("[FlyDSL gather_kv_b_proj] outputs must be bf16")

• Files reviewed: 4/4 changed files
• Comments generated: 0 new
• Review effort level: Lite

Suppressed comments (4)

Previously missed (3) — in code that hasn't changed since the last review.

aiter/ops/flydsl/gather_kv_b_proj.py:199

• gather_kv_b_proj_flydsl relies on kv_indices being int32 and on all tensors living on the same CUDA/ROCm device; otherwise the kernel will reinterpret memory incorrectly or fail inside _run_compiled with a hard-to-diagnose error. Add explicit device/dtype/shape checks for kv_indices, k_scale, outputs, and kv_proj_scale to match validation patterns used in other FlyDSL entrypoints (e.g. gemm_kernels.py).
  op_tests/test_flydsl_gather_kv_b_proj.py:315
• This determinism test allocates a KV cache with num_blocks=3_000_000 (3e6 x 576 fp8 bytes ≈ 1.7GB, plus outputs), which is very likely to OOM on many gfx950 test runners and makes the unit test suite fragile. Consider reducing num_blocks to a smaller value that still exercises the >2^17 address path, or gate it behind an explicit slow/large-memory marker.
  aiter/ops/flydsl/gather_kv_b_proj.py:174
• The kv_proj_scale=None error message claims this backend is "per-output-row scale only", but the implementation also supports 128x128 block scales (see per_row_scale detection and BlockScale). Update the message to reflect the actual supported scale modes so callers aren’t misled.

aiter/ops/flydsl/kernels/gather_kv_b_proj_8wave.py:288

• _rope_copy iterates over ceil(BLOCK_M/256) passes but does not mask out row_local >= block_m. For BLOCK_M values like 384 this makes each workgroup write the first 128 rows of the next tile as well, creating cross-workgroup overlapping stores (data race) and unnecessary traffic. Add a row_local < block_m guard around the index load + copy loop so each workgroup only writes its own tile rows.

    for _p in range_constexpr(n_passes):
        row_local = tid // 2 + _p * ROWS_PER_PASS
        fx.copy(atom32, fx.slice(kvi_div, (None, m_base + row_local)), ireg)
        idx = Vec(fx.memref_load_vec(ireg))[0]
        src_base = idx * KV_ROW_ELEMS + KV_C_DIM + half * 32

• Files reviewed: 4/4 changed files
• Comments generated: 0 new
• Review effort level: Lite

Suppressed comments (2)

Previously missed (1) — in code that hasn't changed since the last review.

aiter/ops/flydsl/gather_kv_b_proj.py:174

• The ValueError for kv_proj_scale=None says the backend is "fp8 x per-output-row scale only", but this backend also supports 128x128 block scaling (kv_proj_scale 2-D). The current message is misleading for users who hit this error with block-scale weights.

aiter/ops/flydsl/kernels/gather_kv_b_proj_8wave.py:333

• _RopeCopy uses row_local = tid//2 + pass*256, but there is no guard to keep row_local < BLOCK_M. For BLOCK_M=128 it will still process rows 128..255, and for BLOCK_M=384 it processes up to 511 (2 passes), causing overlapping stores between adjacent M-tiles (workgroup data race) and extra memory traffic.

        for _p in range_constexpr(self.n_passes):
            row_local = self.tid // 2 + _p * self.ROWS_PER_PASS
            fx.copy(
                self.atom32,
                fx.slice(self.kvi_div, (None, self.m_base + row_local)),

• Files reviewed: 4/4 changed files
• Comments generated: 1
• Review effort level: Lite

Suppressed comments (3)

Previously missed (2) — in code that hasn't changed since the last review.

aiter/ops/flydsl/gather_kv_b_proj.py:321

• k_prefix.view(-1) / v_prefix.view(-1) will throw a low-level PyTorch error if either output is non-contiguous. Since this is a public backend entrypoint that writes outputs in-place, it would be better to validate contiguity explicitly and raise a clear ValueError.
  op_tests/test_flydsl_gather_kv_b_proj.py:86
• For large num_blocks (e.g. the determinism test uses 3,000,000), torch.randperm(num_blocks) is expensive and dominates test setup time/memory even though only alloc entries are used. For large num_blocks, sampling indices with torch.randint is much cheaper and still exercises the gather path.

aiter/ops/flydsl/kernels/gather_kv_b_proj_8wave.py:333

• _RopeCopy maps 512 threads to 256 rows per pass, but load_idx/load_data/commit do not guard row_local < block_m. For BLOCK_M values not divisible by 256 (e.g. 128 or 384), each workgroup will also load/store rope for rows outside its own M tile, overlapping adjacent workgroups’ output rows (redundant writes + cross-workgroup data race, even if values happen to match). Consider masking the rope copy to only the first block_m rows of the tile (and optionally (m_base + row_local) < m_rows to avoid unnecessary OOB loads in the last tile).

        for _p in range_constexpr(self.n_passes):
            row_local = self.tid // 2 + _p * self.ROWS_PER_PASS
            fx.copy(
                self.atom32,
                fx.slice(self.kvi_div, (None, self.m_base + row_local)),

• Files reviewed: 4/4 changed files
• Comments generated: 0 new
• Review effort level: Lite

Suppressed comments (1)

aiter/ops/flydsl/gather_kv_b_proj.py:259

• kv_indices is treated as a packed int32 buffer in the kernel (BufferCopy32b into an fx.Int32 reg). If callers pass int64 (common default), the kernel will reinterpret bytes and silently miscompute; this should be validated up front.

    if kv_indices.numel() < m_rows:
        raise ValueError(
            f"[FlyDSL gather_kv_b_proj] kv_indices has {kv_indices.numel()} entries, "
            f"need at least num_tokens={m_rows}"
        )

• Files reviewed: 4/4 changed files
• Comments generated: 3
• Review effort level: Lite

Suppressed comments (3)

aiter/ops/flydsl/kernels/gather_gemm_8wave.py:355

• _RopeCopy.commit() unconditionally stores 256 rows per pass (based on tid//2), which overwrites the next tile when BLOCK_M=128 (and races other CTAs). Guard the rope stores so threads whose row_local >= BLOCK_M do not write.

        for _p in range_constexpr(self.n_passes):
            for sub in range_constexpr(2):
                words = Vec(fx.memref_load_vec(self.src_regs[_p][sub])).bitcast(
                    fx.Int32
                )

aiter/ops/flydsl/gather_kv_b_proj.py:259

• gather_kv_b_proj_flydsl relies on kv_indices being int32 on the same CUDA device, and on k_scale being a single-element CUDA tensor; otherwise compilation can fail or the kernel can silently read the wrong bytes. Add explicit dtype/device/shape validation here (similar to other ops that assert kv_indices.dtype == torch.int32).

    if kv_indices.numel() < m_rows:
        raise ValueError(
            f"[FlyDSL gather_kv_b_proj] kv_indices has {kv_indices.numel()} entries, "
            f"need at least num_tokens={m_rows}"
        )

aiter/ops/flydsl/gather_kv_b_proj.py:318

• This wrapper flattens outputs via k_prefix.view(-1) / v_prefix.view(-1), which will throw a RuntimeError for non-contiguous outputs. Fail fast with a clear ValueError before launching so callers know contiguity is required for in-place writes.

    _run_compiled(
        exe,
        _as_i8(k_buffer.contiguous()).view(-1),
        kv_indices.contiguous().view(-1),
        _as_i8(kv_proj_weight.contiguous()).view(-1),
        scale.contiguous(),
        k_scale.reshape(-1).to(torch.float32).contiguous(),
        k_prefix.view(-1),
        v_prefix.view(-1),

• Files reviewed: 4/4 changed files
• Comments generated: 1
• Review effort level: Lite

Suppressed comments (4)

Previously missed (1) — in code that hasn't changed since the last review.

aiter/ops/flydsl/gather_kv_b_proj.py:220

• This code assumes all tensors are CUDA and on the same device as k_buffer (e.g. it calls _arch_of(k_buffer.device.index) and passes raw pointers to the compiled kernel). Without explicit device checks, non-CUDA or cross-device inputs can fail with unclear errors or miscompile; kv_indices also needs to be enforced as int32 because the kernel loads indices as 32-bit.

This issue also appears on line 310 of the same file.

aiter/ops/flydsl/kernels/gather_gemm_8wave.py:299

• _RopeCopy needs the runtime BLOCK_M value later (e.g. to mask out row_local >= BLOCK_M in commit), but init does not store block_m anywhere; that makes it hard to prevent redundant/out-of-tile rope stores when BLOCK_M < 256 or not a multiple of 256.

        self.tid = tid
        self.half = tid % 2
        self.n_passes = (block_m + self.ROWS_PER_PASS - 1) // self.ROWS_PER_PASS
        self.kp_row = n_heads * (nope + KV_PE_DIM)
        self.col_base = head * (nope + KV_PE_DIM) + nope

aiter/ops/flydsl/kernels/gather_gemm_8wave.py:355

• _RopeCopy.commit() writes rope data for every row_local in each pass (256 rows/pass), even when the last pass is partial (e.g. BLOCK_M=128/384). That can make multiple blocks store the same output rows (outside their assigned M-tile), creating inter-block overlapping stores (a data race) and extra memory traffic.

        for _p in range_constexpr(self.n_passes):
            for sub in range_constexpr(2):
                words = Vec(fx.memref_load_vec(self.src_regs[_p][sub])).bitcast(
                    fx.Int32
                )

aiter/ops/flydsl/gather_kv_b_proj.py:314

• k_prefix.view(-1) and v_prefix.view(-1) require contiguous storage; if a caller passes a non-contiguous view (e.g. sliced/strided workspace), this will raise a RuntimeError at launch time. Consider validating contiguity explicitly and failing with a clear ValueError.

    _run_compiled(
        exe,
        _as_i8(k_buffer.contiguous()).view(-1),
        kv_indices.contiguous().view(-1),
        _as_i8(kv_proj_weight.contiguous()).view(-1),

• Files reviewed: 4/4 changed files
• Comments generated: 0 new
• Review effort level: Lite